import sys
from pathlib import Path

import whisper
import itertools
import re

RE_STOP_WORDS = re.compile(r'[!?.！？。\s\r\t\n ]+')

root = Path('.').resolve(strict=True)

print(f'当前工作目录：{root}')
model_name = input('要加载的模型名: ')
print('正在加载模型…')
model = whisper.load_model(model_name, in_memory=True, device='cuda')
print('')


def main_loop():
    input_audio_path = input('音频路径 > ')

    glob_path = [x for x in root.glob(f'{input_audio_path}*')]
    if len(glob_path) == 0:
        print('音频文件不存在。\n')
        return

    if len(glob_path) > 1:
        print(f'命中的音频文件数超过 1 个，目前选择的音频：{glob_path[0]}')

    print('正在加载音频…')
    try:
        audio = whisper.load_audio(str(glob_path[0]))
    except Exception as e:
        print('音频加载出错。')
        print(e, file=sys.stderr)
        return

    try:
        print('正在裁剪 / 填充音频…')
        audio = whisper.pad_or_trim(audio)

        print('正在将音频转换为梅尔图…')
        mel = whisper.log_mel_spectrogram(audio)
        audio = None

        print('正在将梅尔图复制到显卡…')
        mel = mel.to(model.device)

        print('模型正在运算…')
        cold_options = whisper.DecodingOptions(
            temperature=0,
            language='Japanese',
            fp16=True,
        )
        result = whisper.decode(model, mel, cold_options)

        print(f'\n冷结果：{result.text}\n')

        history_result = {RE_STOP_WORDS.sub('', result.text)}

        hot_options = whisper.DecodingOptions(
            temperature=1,
            language='Japanese',
            fp16=True,
        )
        for count in itertools.count(1):
            result = whisper.decode(model, mel, hot_options)
            sanitized_result = RE_STOP_WORDS.sub('', result.text)

            if sanitized_result not in history_result:
                history_result.add(sanitized_result)
                print(f'热结果 {count}：{result.text}')

    except KeyboardInterrupt:
        pass

    print('')


while True:
    try:
        main_loop()
    except KeyboardInterrupt:
        exit(0)
